import numpy as np
import cv2
import glob
import argparse

class StereoCalibration:
    def __init__(self, filepath):
        self.boardSize = (11, 8)
        self.squareSize = 1.5  # mm

        self.criteria = (cv2.TERM_CRITERIA_EPS+cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)

        # prepare object points, like (0,0,0), (1,0,0), (2,0,0) ....,(6,5,0)
        self.objp = np.zeros((self.boardSize[0]*self.boardSize[1], 3), np.float32)
        self.objp[:, :2] = np.mgrid[0:self.boardSize[0], 0:self.boardSize[1]].T.reshape(-1, 2)
        self.objp = self.objp * self.squareSize

        self.objpoints = []  # 3d point in real world space
        self.imagePointsLeft = []  # 2d points in image plane.
        self.imagePointsRight = []  # 2d points in image plane.

        self.read_images(filepath)


    def read_images(self, cal_path):

        print(cal_path)
        
        imagePathListLeft = glob.glob(cal_path+"/imageL_*.bmp")
        imagePathListRight = glob.glob(cal_path+"/imageR_*.bmp")
        print(len(imagePathListLeft))

        imagePathListLeft.sort()
        imagePathListRight.sort()

        if len(imagePathListLeft) == len(imagePathListRight):
            image_num = len(imagePathListLeft)

        for i in range(image_num):
            print(i)
            imageLeft = cv2.imread(imagePathListLeft[i])
            imageRight = cv2.imread(imagePathListRight[i])

            imageLeft = cv2.cvtColor(imageLeft, cv2.COLOR_BGR2GRAY)
            imageRight = cv2.cvtColor(imageRight, cv2.COLOR_BGR2GRAY)

            retLeft, cornersLeft = cv2.findChessboardCorners(imageLeft, self.boardSize, None)
            retRight, cornersRight = cv2.findChessboardCorners(imageRight, self.boardSize, None)

            
            if retLeft and retRight:
                cv2.cornerSubPix(imageLeft, cornersLeft, (11, 11), (-1, -1), self.criteria)
                cv2.cornerSubPix(imageRight, cornersRight, (11, 11), (-1, -1), self.criteria)
                self.objpoints.append(self.objp)
                self.imagePointsLeft.append(cornersLeft)
                self.imagePointsRight.append(cornersRight)

            imageShape = imageLeft.shape[::-1]

        _, self.M1, self.D1, self.R1, self.T1 = cv2.calibrateCamera(self.objpoints, self.imagePointsLeft, imageShape, None, None)
        _, self.M2, self.D2, self.R2, self.T2 = cv2.calibrateCamera(self.objpoints, self.imagePointsRight, imageShape, None, None)

        self.stereo_calibrate(imageShape)

    def stereo_calibrate(self, dims):
        flags = 0
        flags |= cv2.CALIB_FIX_INTRINSIC
        # flags |= cv2.CALIB_FIX_PRINCIPAL_POINT
        flags |= cv2.CALIB_USE_INTRINSIC_GUESS
        flags |= cv2.CALIB_FIX_FOCAL_LENGTH
        # flags |= cv2.CALIB_FIX_ASPECT_RATIO
        flags |= cv2.CALIB_ZERO_TANGENT_DIST
        # flags |= cv2.CALIB_RATIONAL_MODEL
        # flags |= cv2.CALIB_SAME_FOCAL_LENGTH
        # flags |= cv2.CALIB_FIX_K3
        # flags |= cv2.CALIB_FIX_K4
        # flags |= cv2.CALIB_FIX_K5

        stereocalib_criteria = (cv2.TERM_CRITERIA_MAX_ITER +
                                cv2.TERM_CRITERIA_EPS, 100, 1e-5)
        ret, M1, D1, M2, D2, R, T, E, F = cv2.stereoCalibrate(
            self.objpoints, self.imagePointsLeft,
            self.imagePointsRight, self.M1, self.D1, self.M2,
            self.D2, dims,
            criteria=stereocalib_criteria, flags=flags)
        print('ret', ret)
        print('Intrinsic_mtx_1', M1)
        print('dist_1', D1)
        print('Intrinsic_mtx_2', M2)
        print('dist_2', D2)
        print('R', R)
        print('T', T)
        print('E', E)
        print('F', F)

        # for i in range(len(self.r1)):
        #     print("--- pose[", i+1, "] ---")
        #     self.ext1, _ = cv2.Rodrigues(self.r1[i])
        #     self.ext2, _ = cv2.Rodrigues(self.r2[i])
        #     print('Ext1', self.ext1)
        #     print('Ext2', self.ext2)

        print('')

        camera_model = dict([('M1', M1), ('M2', M2), ('dist1', D1),
                            ('dist2', D2), ('rvecs1', self.R1),
                            ('rvecs2', self.R2), ('R', R), ('T', T),
                            ('E', E), ('F', F)])

        # cv2.destroyAllWindows()
        return camera_model
    
if __name__ == '__main__':
    
    cal_data = StereoCalibration("/mnt/d/data/CameraCalibration")


    # ret 0.4377717434408196
    # Intrinsic_mtx_1 [[1.80373929e+04 0.00000000e+00 1.09501461e+03]
    # [0.00000000e+00 1.82138154e+04 1.06707196e+03]
    # [0.00000000e+00 0.00000000e+00 1.00000000e+00]]
    # dist_1 [[-1.01022863e+00  1.52268137e+02  3.06688761e-03 -1.61819434e-02
    # 5.25883387e-01]]
    # Intrinsic_mtx_2 [[1.79091679e+04 0.00000000e+00 9.82910335e+02]
    # [0.00000000e+00 1.78168563e+04 1.65093488e+03]
    # [0.00000000e+00 0.00000000e+00 1.00000000e+00]]
    # dist_2 [[ 2.07600160e-01 -1.25861492e+02  2.15050575e-03  6.39670396e-03
    # 4.60303554e-02]]
    # R [[ 0.6243645  -0.00246236  0.78112925]
    # [ 0.02592041  0.99950963 -0.01756768]
    # [-0.78070296  0.03121583  0.62412216]]
    # T [[-187.00489078]
    # [  -3.22129786]
    # [  85.80532326]]
    # E [[   0.29076766  -85.86380274   -0.50308297]
    # [ -92.42147316    5.62622816  183.73894397]
    # [  -2.83597928 -186.921122      5.80149193]]
    # F [[-1.07440855e-09  3.14200096e-07 -3.00567372e-04]
    # [ 3.43273737e-07 -2.06946396e-08 -1.26633712e-02]
    # [-3.77993529e-04  1.19751628e-02  1.00000000e+00]]
